package org.iscas.tcse.datasample;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.ml.linalg.DenseVector;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import redis.clients.jedis.Jedis;

import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;

public class RandomSampler extends ProcessFunction<DenseVector, DenseVector> implements Sampler {

    private static final Logger LOG = LoggerFactory.getLogger(RandomSampler.class);
    private volatile double sampleRatio;

    private volatile AtomicInteger counter;

    private final Random random;
    private final SamplerConf conf;

    private SampleController controller = null;

    private final boolean enableController;

    private volatile boolean isRunning = true;

    public RandomSampler(double sampleRatio, SamplerConf conf, boolean enableController) {
        this.random = new Random(1024);
        this.sampleRatio = sampleRatio;
        this.conf = conf;
        this.enableController = enableController;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        if (enableController) {
            int taskId = getRuntimeContext().getIndexOfThisSubtask();
            if (taskId == 0) {
                Thread controller = new Thread(new SampleController(this.conf));
                controller.start();
                this.controller = (SampleController) controller;
                this.counter = new AtomicInteger(0);
                new Thread(
                                () -> {
                                    try {
                                        reportSourceSpeed();
                                    } catch (InterruptedException e) {
                                        e.printStackTrace();
                                    }
                                })
                        .start();
            }
            new Thread(
                            () -> {
                                try {
                                    syncSampleRate();
                                } catch (InterruptedException e) {
                                    e.printStackTrace();
                                }
                            })
                    .start();
        }
    }

    @Override
    public void close() throws Exception {
        super.close();
        if (this.controller != null) {
            this.controller.stopRunning();
        }
        this.isRunning = false;
    }

    @Override
    public void processElement(
            DenseVector value,
            ProcessFunction<DenseVector, DenseVector>.Context ctx,
            Collector<DenseVector> out)
            throws Exception {
        boolean shouldRelease = (random.nextDouble() < sampleRatio);
        if (shouldRelease) {
            out.collect(value);
        }
        counter.getAndIncrement();
    }

    @Override
    public void reportSourceSpeed() throws InterruptedException {
        Jedis jedis = new Jedis(conf.redisHost, conf.redisPort);
        while (isRunning) {
            jedis.set(conf.sourceSpeedKey, String.valueOf(counter));
            counter.set(0);
            Thread.sleep(conf.syncInterval);
        }
        jedis.close();
    }

    public void syncSampleRate() throws InterruptedException {
        Jedis jedis = new Jedis(conf.redisHost, conf.redisPort);
        while (isRunning) {
            String sampleRatioStr = jedis.get(conf.sampleRatioKey);
            if (sampleRatioStr != null) {
                this.sampleRatio = Double.parseDouble(sampleRatioStr);
            }
            Thread.sleep(conf.syncInterval);
        }
        jedis.close();
    }
}
